"""
EEG信号处理的图神经网络实现
包含数据预处理、图构建、模型训练和评估等功能
"""

import torch
import torch.nn as nn


class AutomaticWeightedLoss(nn.Module):
    """
    自动加权多任务损失函数
    通过可学习参数自动调整不同损失项的权重

    参数:
        num: 需要加权的损失项数量
    """

    def __init__(self, num=2):
        super(AutomaticWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)  # 初始化权重参数
        self.params = torch.nn.Parameter(params)  # 注册为可学习参数
        print(self.params)

    def forward(self, *x):
        """
        前向计算加权损失

        参数:
            *x: 可变参数，传入多个损失项
        返回:
            loss_sum: 加权后的总损失
        """
        loss_sum = 0
        length = len(x) - 1
        for i, loss in enumerate(x):
            # 对最后一个损失项使用不同的加权方式
            if i == length:
                loss_sum += 1 / (self.params[i] ** 2) * loss + torch.log(self.params[i])
            else:
                loss_sum += 0.5 / (self.params[i] ** 2) * loss + torch.log(self.params[i])
        return loss_sum


# 主程序入口
if __name__ == '__main__':
    awl = AutomaticWeightedLoss(2)  # 测试自动加权损失函数
    print(awl.parameters())
